#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件        :    make_yolo_mini.py
@说明        :    
@时间        :    2021/12/21 18:59:26
@作者        :    钟建平
@版本        :    1.0
'''

import os
import argparse
import numpy as np
import shutil
from tqdm import tqdm

def  _parser_args():
    """parse input"""
    parser = argparse.ArgumentParser("convert coco to coco mini")
    parser.add_argument('--yolo_path', type=str, required=True, help='the path of coco')
    parser.add_argument('--ratio', type=float, required=True, help='the path of coco')
    parser.add_argument('--yolo_mini_path', type=str,  required=True,  help='the path ot save coco_mini data')
    args, _ = parser.parse_known_args()
    return args


def split(data_type, ratio, src_path, dst_path):
    print(f'---------------start to process {data_type}--------------------')
    if not os.path.exists(dst_path+f'images/{data_type}'):
        os.mkdir(dst_path+f'images/{data_type}')
    if not os.path.exists(dst_path+f'labels/{data_type}'):
        os.mkdir(dst_path+f'labels/{data_type}')
    label_list = os.listdir(src_path+f'labels/{data_type}')
    np.random.shuffle(label_list)
    for j in tqdm(range(int((len(label_list)*ratio)))):
        label_name = f'labels/{data_type}/{label_list[j]}'
        image_name = f"images/{data_type}/{label_list[j].replace('.txt', '.jpg')}"
        shutil.copy(src_path+image_name, dst_path+image_name)
        shutil.copy(src_path+label_name, dst_path+label_name)
    print('-------------process {data_type} finished!-------------------------')
def make(yolo_path, ratio, yolo_mini_path):
    if not os.path.exists(yolo_mini_path):
        os.mkdir(yolo_mini_path)
    if not os.path.exists(yolo_mini_path+'images'):
        os.mkdir(yolo_mini_path+'images')
    if not os.path.exists(yolo_mini_path+'labels'):
        os.mkdir(yolo_mini_path+'labels')
    split('train2017', ratio, yolo_path, yolo_mini_path)
    split('val2017', ratio,yolo_path, yolo_mini_path)

if __name__ == '__main__':
    args = _parser_args()
    make(args.yolo_path, args.ratio, args.yolo_mini_path)
